import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchsummary import summary

# 与SE不同的点在于 注意力机制从全局平均池化GAP变为二维离散余弦变换DCT
# 作者证明了GAP是DCT的特殊形式，因此使用DCT可以更加突出通道的特征信息
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# [u,v]表示二维DCT的分量下标，即实现对每一块采用不同的频率分量
# 这里u,v的值是根据论文中所做实验,选取正确率top_k得出
# 因为初始化的固定张量效果最好，所以这么写,开源代码里取top32,这里取top16
fidx_u_16 = [0,0,6,0,0,1,1,4,5,1,3,0,0,0,2,3] # 论文所选频率的水平指数
fidx_v_16 = [0,1,0,5,2,0,2,0,0,6,0,4,6,3,2,5] # 论文所选频率的垂直指数
# 获取一维通道的DCT
# L为输入的长度
def get_1d_dct(pos, freq, L):
    result = math.cos(math.pi * freq * (pos + 0.5) / L) / math.sqrt(L)
    if freq == 0:
        return result
    else:
        return result * math.sqrt(2)

# 推广到二维DCT
def get_dct_weights(input_width, input_height, input_channel, fidx_u= fidx_u_16, fidx_v= fidx_v_16):
    # 论文将频率空间分为7*7,实验结果表明,[u,v]=[0,0],即GAP的效果最好(乐),当然了，改进的点在于在信道注意力中使用多个频率分量
    # 使不同大小的频率与 7x7 频率空间相同
    fidx_horizon = [u*(input_width//7) for u in fidx_u]
    fidx_vertical = [v*(input_height//7) for v in fidx_v]
    dct_weights = torch.zeros(1, input_channel, input_width, input_height)
    c_part = input_channel // len(fidx_horizon)
    # 拆分通道
    for i, (u_x, v_y) in enumerate(zip(fidx_horizon, fidx_vertical)):
        for t_x in range(input_width):
            for t_y in range(input_height):
                dct_weights[:, i * c_part: (i+1)*c_part, t_x, t_y]\
                =get_1d_dct(t_x, u_x, input_width) * get_1d_dct(t_y, v_y, input_height)

    return dct_weights

# SE-Block单元--SEblock是一个子结构，几乎可以嵌入任何一个神经网络模型之中
class FcaBlock(nn.Module):
    def __init__(self, input_channel, width, height, reduction=16):
        super(FcaBlock, self).__init__()
        self.width = width
        self.height = height


        # self.adaptive_avg_pool = nn.AdaptiveAvgPool2d(1)  # 全局自适应池化
        self.register_buffer('pre_computed_dct_weights', get_dct_weights(self.width, self.height, input_channel))
        self.fc = nn.Sequential(
            nn.Linear(input_channel, input_channel // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(input_channel // reduction, input_channel),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, h, w = x.size()

        # squeeze操作:(b,c,h,w)->(b,c)
        #y = self.adaptive_avg_pool(x).view(b, c)
        y = F.adaptive_avg_pool2d(x, (self.height, self.width))
        y = torch.sum(y*self.pre_computed_dct_weights, dim=[2,3])

        # FC获取通道注意力权重，是具有全局信息的
        y = self.fc(y).view(b, c, 1, 1)

        # 注意力作用每一个通道上
        y = x * y.expand_as(x)
        # 残差连接
        return x+y


class Attention(nn.Module):
    def __init__(self, in_channels, out_channels, width, height, reduction=16, rate=4):
        super(Attention, self).__init__()
        self.width = width
        self.height = height

        self.spatial_attention = nn.Sequential(
            nn.Conv2d(in_channels, int(in_channels / rate), kernel_size=3, padding=1),
            nn.BatchNorm2d(int(in_channels / rate)),
            nn.ReLU(inplace=True),
            nn.Conv2d(int(in_channels / rate), out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels)
        )

        self.register_buffer('pre_computed_dct_weights', get_dct_weights(self.width, self.height, in_channels))
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels),
            nn.Sigmoid()
        )
    def forward(self, x):
        b, c, h, w = x.shape
        x_spatial_att = self.spatial_attention(x).sigmoid()
        x = x * x_spatial_att

        y = F.adaptive_avg_pool2d(x, (self.height, self.width))
        y = torch.sum(y * self.pre_computed_dct_weights, dim=[2, 3])
        y = self.fc(y).view(b, c, 1, 1)

        y = x * y.expand_as(x)

        return x + y

# 用于Resnet18 Resnet34
class FcaBasicBlock(nn.Module):
    expansion = 1

    def __init__(self, in_channels, out_channels, stride=1, reduction=16):
        super().__init__()
        # c2wh = dict([(64, 56), (128, 28), (256, 14), (512, 7)])
        c2wh = dict([(64, 28), (128, 14), (256, 7), (512, 4)])
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(out_channels, out_channels * FcaBasicBlock.expansion, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels * FcaBasicBlock.expansion)

        #self.se = FcaBlock(out_channels, reduction)
        # self.fca = FcaBlock(out_channels, c2wh[out_channels], c2wh[out_channels], reduction)
        self.att = Attention(out_channels, out_channels, c2wh[out_channels], c2wh[out_channels], reduction)

        self.residual = nn.Sequential()


        if stride != 1 or in_channels != FcaBasicBlock.expansion * out_channels:
            self.residual = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * FcaBasicBlock.expansion, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels * FcaBasicBlock.expansion)
            )

    def forward(self, x):
        residual = self.residual(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)

        out = self.att(out)
        out = out * out

        out = out + residual
        out = F.relu(out)
        return out

class FcaBottleNeck(nn.Module):
    expansion = 4
    def __init__(self, in_channels, out_channels, stride=1, reduction=16):
        super().__init__()
        # c2wh = dict([(64, 56), (128, 28), (256, 14), (512, 7)])
        c2wh = dict([(64, 28), (128, 14), (256, 7), (512, 4)])
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

        # 第二个卷积层：卷积核尺寸： 3×3 ，填充值为 1， 步长为 1
        self.conv2 = nn.Conv2d(out_channels, out_channels, stride=stride, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)

        # 第三个卷积层：卷积核尺寸： 1×1 ，填充值为 0， 步长为 1
        self.conv3 = nn.Conv2d(out_channels, out_channels * FcaBottleNeck.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels * FcaBottleNeck.expansion)
        self.relu = nn.ReLU(inplace=True)

        #self.se = FcaBlock(out_channels * SEBottleNeck.expansion, reduction)
        # self.fca = FcaBlock(out_channels*FcaBottleNeck.expansion, c2wh[out_channels], c2wh[out_channels], reduction)
        self.att = Attention(out_channels * FcaBottleNeck.expansion, out_channels * FcaBottleNeck.expansion, c2wh[out_channels], c2wh[out_channels], reduction)

        self.residual = nn.Sequential()

        if stride != 1 or in_channels != out_channels * FcaBottleNeck.expansion:
            self.residual = nn.Sequential(
                nn.Conv2d(in_channels, out_channels * FcaBottleNeck.expansion, stride=stride, kernel_size=1, bias=False),
                nn.BatchNorm2d(out_channels * FcaBottleNeck.expansion)
            )

    def forward(self, x):
        residual = self.residual(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        out = self.conv3(out)
        out = self.bn3(out)
        out = self.relu(out)
        out = self.att(out)
        out = out * out

        out = out + residual
        out = F.relu(out)
        return out

class FcaResNet(nn.Module):

    def __init__(self, block, num_block, num_classes=100, dropout_prob0=0.3):
        super().__init__()

        self.in_channels = 64

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True))

        self.conv2_x = self._make_layer(block, 64, num_block[0], 1)
        self.conv3_x = self._make_layer(block, 128, num_block[1], 2)
        self.conv4_x = self._make_layer(block, 256, num_block[2], 2)
        self.conv5_x = self._make_layer(block, 512, num_block[3], 2)
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))

        if dropout_prob0 > 0.0:
            self.dp = nn.Dropout(dropout_prob0, inplace=True)
        else:
            self.dp = None

        self.fc = nn.Linear(512 * block.expansion, num_classes)

    def _make_layer(self, block, out_channels, num_blocks, stride):

        strides = [stride] + [1] * (num_blocks - 1)
        layers = []
        for stride in strides:
            layers.append(block(self.in_channels, out_channels, stride))
            self.in_channels = out_channels * block.expansion

        return nn.Sequential(*layers)

    def forward(self, x):
        output = self.conv1(x)
        output = self.conv2_x(output)
        output = self.conv3_x(output)
        output = self.conv4_x(output)
        output = self.conv5_x(output)
        output = self.avg_pool(output)
        output = output.view(output.size(0), -1)
        output = self.fc(output)

        return output
